{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: linear evolution with re-linearisation during evolutive equilibrium calculations\n",
    "\n",
    "This notebook also follows on from the previous two. \n",
    "\n",
    "Here, we demonstrate how to re-calculate the linearisation (e.g. Jacobians) of the dynamics (\"relinearise\") when using the linear solver. As we know, the accuracy of the Jacobians (calculated using the initial equilibrium, before evolution) will degrade as the plasma \"moves away\" from the initial equilibrium around which it was calculated. This means that we need to re-linearise every so often in order to maintain accuracy of the plasma evolution over time. \n",
    "\n",
    "The more often a relinearisation occurs, the more accurate (in theory) a linear simulation will be (when compared to the full nonlinear simulation). It is not practical or efficient to relinearise each timestep, but doing so once a given criterion is met can significantly improve agreement with the nonlinear solver. To trigger relinearisation, we monitor the relative change in plasma current density $J_{\\phi}$ at the current time $t$ and the initial time $t_0$. If this threshold exceeds tolerance $\\epsilon$, then we re-linearise around the equilibrium at time $t$ (and set $t_0 = t$). \n",
    "\n",
    "The criterion checks whether\n",
    "$$\n",
    "\\frac{\\left \\lVert J_{\\phi}(t) - J_{\\phi}(t_0) \\right \\rVert}{\\left \\lVert J_{\\phi}(t_0) \\right \\rVert} < \\epsilon,\n",
    "$$\n",
    "is met or not at each timestep $t$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display, clear_output\n",
    "import pickle\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# build machine\n",
    "from freegsnke import build_machine\n",
    "tokamak = build_machine.tokamak(\n",
    "    active_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_active_coils.pickle\",\n",
    "    passive_coils_path=f\"../machine_configs/MAST-U/MAST-U_like_passive_coils.pickle\",\n",
    "    limiter_path=f\"../machine_configs/MAST-U/MAST-U_like_limiter.pickle\",\n",
    "    wall_path=f\"../machine_configs/MAST-U/MAST-U_like_wall.pickle\",\n",
    ")\n",
    "\n",
    "# initialise equilibrium object\n",
    "from freegsnke import equilibrium_update\n",
    "eq = equilibrium_update.Equilibrium(\n",
    "    tokamak=tokamak,\n",
    "    Rmin=0.1, Rmax=2.0,   # radial range\n",
    "    Zmin=-2.2, Zmax=2.2,  # vertical range\n",
    "    nx=65,                # number of grid points in the radial direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    ny=129,               # number of grid points in the vertical direction (needs to be of the form (2**n + 1) with n being an integer)\n",
    "    # psi=plasma_psi\n",
    ")  \n",
    "\n",
    "# initialise profile object\n",
    "from freegsnke.jtor_update import ConstrainPaxisIp\n",
    "profiles = ConstrainPaxisIp(\n",
    "    eq=eq,\n",
    "    paxis=8.1e3,\n",
    "    Ip=6.2e5,\n",
    "    fvac=0.5,\n",
    "    alpha_m=1.8,\n",
    "    alpha_n=1.2\n",
    ")\n",
    "\n",
    "# initialise solver\n",
    "from freegsnke import GSstaticsolver\n",
    "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)    \n",
    "\n",
    "# set coil currents\n",
    "import pickle\n",
    "with open('data/simple_diverted_currents_PaxisIp.pk', 'rb') as f:\n",
    "    current_values = pickle.load(f)\n",
    "\n",
    "for key in current_values.keys():\n",
    "    eq.tokamak[key].current = current_values[key]\n",
    "\n",
    "# carry out forward solve\n",
    "GSStaticSolver.solve(eq=eq, \n",
    "                     profiles=profiles, \n",
    "                     constrain=None, \n",
    "                     target_relative_tolerance=1e-9)\n",
    "\n",
    "# plot the resulting equilbrium\n",
    "fig1, ax1 = plt.subplots(1, 1, figsize=(4, 8), dpi=80)\n",
    "ax1.grid(True, which='both')\n",
    "eq.plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the descriptors function (it should return an array of values and take in an eq object)\n",
    "def plasma_descriptors(eq):\n",
    "\n",
    "    # find lower X-point\n",
    "    # define a \"box\" in which to search for the lower X-point\n",
    "    XPT_BOX = [[0.33, -0.88], [0.95, -1.38]]\n",
    "\n",
    "    # mask those points\n",
    "    xpt_mask = (\n",
    "        (eq.xpt[:, 0] >= XPT_BOX[0][0])\n",
    "        & (eq.xpt[:, 0] <= XPT_BOX[1][0])\n",
    "        & (eq.xpt[:, 1] <= XPT_BOX[0][1])\n",
    "        & (eq.xpt[:, 1] >= XPT_BOX[1][1])\n",
    "    )\n",
    "    xpts = eq.xpt[xpt_mask, 0:2].squeeze()\n",
    "    if xpts.ndim > 1 and xpts.shape[0] > 1:\n",
    "        opt = eq.opt[0, 0:2]\n",
    "        dists = np.linalg.norm(xpts - opt, axis=1)\n",
    "        idx = np.argmin(dists)  # index of closest point\n",
    "        Rx, Zx = xpts[idx, :]\n",
    "    else:\n",
    "        Rx, Zx = xpts\n",
    "\n",
    "    # find avg. Z position of jtor\n",
    "    Zcurrent = eq.Zcurrent()\n",
    "\n",
    "    # inboard midplane radius\n",
    "    Rin = eq.innerOuterSeparatrix()[0]\n",
    "\n",
    "    return np.array([Zcurrent, Rx, Zx, Rin])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise the nonlinear solver object\n",
    "from freegsnke import nonlinear_solve\n",
    "\n",
    "stepping = nonlinear_solve.nl_solver(\n",
    "    eq=eq, \n",
    "    profiles=profiles, \n",
    "    GSStaticSolver=GSStaticSolver,\n",
    "    full_timestep=5e-4, \n",
    "    plasma_resistivity=1e-6,\n",
    "    max_mode_frequency=10**2.5,\n",
    "    n_linearization_workers=1,  # set >1 to parallelise Jacobian construction\n",
    "    plasma_descriptor_function=plasma_descriptors\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The new **`n_linearization_workers`** argument controls how many independent worker processes build both the current-response and profile-parameter-response columns. It applies to the first linearisation constructed by `nl_solver` and to later relinearisations. The default value of `1` retains the serial calculation. Values larger than one use that many CPU threads in total for Jacobian construction, with one native numerical-library thread per worker to avoid oversubscription. This limit is scoped to Jacobian construction; it does not alter the inner threading used by the ordinary GS solves that advance the evolution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Carry out standard linear solve\n",
    "\n",
    "As we have done before, we carry out a standard linear simulation so that we can compare how the relinearisation performs compared to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise with the initial condition\n",
    "stepping.initialize_from_ICs(eq, profiles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of time steps to simulate\n",
    "max_count = 50\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times = [t]\n",
    "history_currents = [stepping.currents_vec]\n",
    "history_equilibria = [stepping.eq1.create_auxiliary_equilibrium()]\n",
    "history_o_points = [stepping.eq1.opt[0]]\n",
    "history_elongation = [stepping.eq1.geometricElongation()]\n",
    "history_triangularity = [stepping.eq1.triangularity()]\n",
    "history_squareness = [stepping.eq1.squareness()[1]]\n",
    "history_area = [stepping.eq1.separatrix_area()]\n",
    "history_length = [stepping.eq1.separatrix_length()]\n",
    "history_jtor = [stepping.profiles1.jtor]\n",
    "history_jtor_norm = []\n",
    "history_timings = []\n",
    "history_plasma_descriptors = [plasma_descriptors(eq)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use constant voltages\n",
    "voltages = (stepping.vessel_currents_vec*stepping.evol_metal_curr.coil_resist)[:stepping.evol_metal_curr.n_active_coils] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the solver (linear)\n",
    "We call the linear solver without performing any re-linearisations. This will use the linearisation (calculated when `initialize_from_ICs` is called) to evolve the plasma for all iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialise the solver with the initial equilibrium/profiles\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# loop over time steps\n",
    "while counter<max_count:\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "\n",
    "    # start timer\n",
    "    start_time = time.time()\n",
    "\n",
    "    # carry out the time step (profile parameters and resistivity are constant at each time step)\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages,   # same voltages used at each time step\n",
    "        linear_only=True,\n",
    "        verbose=False,\n",
    "    )\n",
    "\n",
    "    # stop timer\n",
    "    end_time = time.time()\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times.append(t)\n",
    "    counter += 1\n",
    "\n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents.append(stepping.currents_vec)\n",
    "    history_equilibria.append(stepping.eq1.create_auxiliary_equilibrium())\n",
    "    history_o_points.append(stepping.eq1.opt[0])\n",
    "    history_elongation.append(stepping.eq1.geometricElongation())\n",
    "    history_triangularity.append(stepping.eq1.triangularity())\n",
    "    history_squareness.append(stepping.eq1.squareness()[1])\n",
    "    history_area.append(stepping.eq1.separatrix_area())\n",
    "    history_length.append(stepping.eq1.separatrix_length())\n",
    "    history_jtor.append(stepping.profiles1.jtor)\n",
    "    history_timings.append(end_time - start_time)\n",
    "    history_jtor_norm.append(np.linalg.norm(stepping.profiles1.jtor - stepping.jtor0) / np.linalg.norm(stepping.jtor0))\n",
    "    history_plasma_descriptors.append(plasma_descriptors(stepping.eq1))\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents = np.array(history_currents)\n",
    "history_times = np.array(history_times)\n",
    "history_o_points = np.array(history_o_points)\n",
    "history_elongation = np.array(history_elongation)\n",
    "history_triangularity = np.array(history_triangularity)\n",
    "history_squareness = np.array(history_squareness)\n",
    "history_area = np.array(history_area)\n",
    "history_length = np.array(history_length)\n",
    "history_jtor = np.array(history_jtor)\n",
    "history_timings = np.array(history_timings)\n",
    "history_jtor_norm = np.array(history_jtor_norm)\n",
    "history_plasma_descriptors = np.array(history_plasma_descriptors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the solver (relinearisation)\n",
    "\n",
    "Again, we evolve the plasma using the linear solver except this time we will provide `relinearise_threshold` to the `nlstepper` call (this is the value of $\\epsilon$ we desire). This will re-calculate the linearisation of the dynamics when the relative change in $J_{\\phi}$ is greater than the threshold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose threshold at which to relinearise automatically (this does it at 5%)\n",
    "threshold = 0.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.dIydI_ICs = None\n",
    "stepping.dIydtheta_ICs = None\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times_rl = [t]\n",
    "history_currents_rl = [stepping.currents_vec]\n",
    "history_equilibria_rl = [stepping.eq1.create_auxiliary_equilibrium()]\n",
    "history_o_points_rl = [stepping.eq1.opt[0]]\n",
    "history_elongation_rl = [stepping.eq1.geometricElongation()]\n",
    "history_triangularity_rl = [stepping.eq1.triangularity()]\n",
    "history_squareness_rl = [stepping.eq1.squareness()[1]]\n",
    "history_area_rl = [stepping.eq1.separatrix_area()]\n",
    "history_length_rl = [stepping.eq1.separatrix_length()]\n",
    "history_timings_rl = []\n",
    "history_jtor_norm_rl = []\n",
    "history_plasma_descriptors_rl = [plasma_descriptors(eq)]\n",
    "\n",
    "# loop over the time steps\n",
    "while counter<max_count:\n",
    "\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "    \n",
    "    # start timer\n",
    "    start_time = time.time()\n",
    "\n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages, \n",
    "        linear_only=True,\n",
    "        verbose=False,\n",
    "        relinearise_threshold=threshold,\n",
    "    )\n",
    "\n",
    "    # stop timer\n",
    "    end_time = time.time()\n",
    "\n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times_rl.append(t)\n",
    "    counter += 1\n",
    "    \n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents_rl.append(stepping.currents_vec)\n",
    "    history_equilibria_rl.append(stepping.eq1.create_auxiliary_equilibrium())\n",
    "    history_o_points_rl.append(stepping.eq1.opt[0])\n",
    "    history_elongation_rl.append(stepping.eq1.geometricElongation())\n",
    "    history_triangularity_rl.append(stepping.eq1.triangularity())\n",
    "    history_squareness_rl.append(stepping.eq1.squareness()[1])\n",
    "    history_area_rl.append(stepping.eq1.separatrix_area())\n",
    "    history_length_rl.append(stepping.eq1.separatrix_length())\n",
    "    history_timings_rl.append(end_time - start_time)\n",
    "    history_jtor_norm_rl.append(stepping.relinearise_criteria)\n",
    "    history_plasma_descriptors_rl.append(plasma_descriptors(stepping.eq1))\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_rl = np.array(history_currents_rl)\n",
    "history_times_rl = np.array(history_times_rl)\n",
    "history_o_points_rl = np.array(history_o_points_rl)\n",
    "history_elongation_rl = np.array(history_elongation_rl)\n",
    "history_triangularity_rl = np.array(history_triangularity_rl)\n",
    "history_squareness_rl = np.array(history_squareness_rl)\n",
    "history_area_rl = np.array(history_area_rl)\n",
    "history_length_rl = np.array(history_length_rl)\n",
    "history_timings_rl = np.array(history_timings_rl)\n",
    "history_jtor_norm_rl = np.array(history_jtor_norm_rl)\n",
    "history_plasma_descriptors_rl = np.array(history_plasma_descriptors_rl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Call the solver (nonlinear)\n",
    "\n",
    "Let us also do a nonlinear simulation to compare the re-linearisation simulation to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.dIydI_ICs = None\n",
    "stepping.dIydtheta_ICs = None\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "# initialising some variables for iteration and logging\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "history_times_nl = [t]\n",
    "history_currents_nl = [stepping.currents_vec]\n",
    "history_equilibria_nl = [stepping.eq1.create_auxiliary_equilibrium()]\n",
    "history_o_points_nl = [stepping.eq1.opt[0]]\n",
    "history_elongation_nl = [stepping.eq1.geometricElongation()]\n",
    "history_triangularity_nl = [stepping.eq1.triangularity()]\n",
    "history_squareness_nl = [stepping.eq1.squareness()[1]]\n",
    "history_area_nl = [stepping.eq1.separatrix_area()]\n",
    "history_length_nl = [stepping.eq1.separatrix_length()]\n",
    "history_timings_nl = []\n",
    "history_jtor_norm_nl = []\n",
    "history_plasma_descriptors_nl = [plasma_descriptors(eq)]\n",
    "\n",
    "# loop over the time steps\n",
    "while counter<max_count:\n",
    "\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "\n",
    "    # start timer\n",
    "    start_time = time.time()\n",
    "\n",
    "    # carry out the time step\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages, \n",
    "        linear_only=False, \n",
    "        verbose=False,\n",
    "        max_solving_iterations=50,\n",
    "        target_relative_tol_currents=1e-2,                # relative tolerance in the currents required for convergence\n",
    "        target_relative_tol_GS=1e-2,                      # relative tolerance in the plasma flux required for convergence\n",
    "        working_relative_tol_GS=5e-3,                     # tolerance used when solving all static GS problems, expressed in terms of the change in the plasma flux due to 1 timestep of evolution\n",
    "    )\n",
    "\n",
    "    # stop timer\n",
    "    end_time = time.time()\n",
    "    \n",
    "    # store information on the time step\n",
    "    t += stepping.dt_step\n",
    "    history_times_nl.append(t)\n",
    "    counter += 1\n",
    "    \n",
    "    # store time-advanced equilibrium, currents, and profiles (+ other quantites of interest)\n",
    "    history_currents_nl.append(stepping.currents_vec)\n",
    "    history_equilibria_nl.append(stepping.eq1.create_auxiliary_equilibrium())\n",
    "    history_o_points_nl.append(stepping.eq1.opt[0])\n",
    "    history_elongation_nl.append(stepping.eq1.geometricElongation())\n",
    "    history_triangularity_nl.append(stepping.eq1.triangularity())\n",
    "    history_squareness_nl.append(stepping.eq1.squareness()[1])\n",
    "    history_area_nl.append(stepping.eq1.separatrix_area())\n",
    "    history_length_nl.append(stepping.eq1.separatrix_length())\n",
    "    history_timings_nl.append(end_time - start_time)\n",
    "    history_jtor_norm_nl.append(stepping.relinearise_criteria)\n",
    "    history_plasma_descriptors_nl.append(plasma_descriptors(stepping.eq1))\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_nl = np.array(history_currents_nl)\n",
    "history_times_nl = np.array(history_times_nl)\n",
    "history_o_points_nl = np.array(history_o_points_nl)\n",
    "history_elongation_nl = np.array(history_elongation_nl)\n",
    "history_triangularity_nl = np.array(history_triangularity_nl)\n",
    "history_squareness_nl = np.array(history_squareness_nl)\n",
    "history_area_nl = np.array(history_area_nl)\n",
    "history_length_nl = np.array(history_length_nl)\n",
    "history_timings_nl = np.array(history_timings_nl)\n",
    "history_jtor_norm_nl = np.array(history_jtor_norm_nl)\n",
    "history_plasma_descriptors_nl = np.array(history_plasma_descriptors_nl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Comparing the simulations\n",
    "We can now plot a comparison of the nonlinear, linear, and linear (with relinearisation) simulations. We also plot the times at which relinearisation took place (as dashed lines). \n",
    "\n",
    "We observe performance of the relinearised simulation somewhere \"inbetween\" the nonlinear and linear simulations as expected! Note that the \"jumps\" are in the triangularity and the squareness parameters are mostly due to resolution effects. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot evolution of tracked values and compare between linear and non-linear evolution\n",
    "fig, axs = plt.subplots(3, 3, figsize=(15, 10), dpi=80, constrained_layout=True)\n",
    "axs_flat = axs.flat\n",
    "\n",
    "axs_flat[0].plot(history_times, history_o_points[:, 0],'k+', label='linear')\n",
    "axs_flat[0].plot(history_times_rl, history_o_points_rl[:, 0],'b.', label='relinearised')\n",
    "axs_flat[0].plot(history_times_nl, history_o_points_nl[:, 0],'rx', label='nonlinear')\n",
    "axs_flat[0].set_xlabel('Time')\n",
    "axs_flat[0].set_ylabel('Magnetic axis, $R$ [$m$]')\n",
    "axs_flat[0].legend()\n",
    "axs_flat[0].grid()\n",
    "\n",
    "\n",
    "axs_flat[1].plot(history_times, abs(history_o_points[:, 1]),'k+')\n",
    "axs_flat[1].plot(history_times_rl, abs(history_o_points_rl[:, 1]),'b.')\n",
    "axs_flat[1].plot(history_times_nl, abs(history_o_points_nl[:, 1]),'rx')\n",
    "axs_flat[1].set_yscale('log')\n",
    "axs_flat[1].set_xlabel('Time')\n",
    "axs_flat[1].set_ylabel('abs(Magnetic axis, $Z$) [$m$]')\n",
    "axs_flat[1].grid()\n",
    "\n",
    "axs_flat[2].plot(history_times, history_o_points[:, 2],'k+')\n",
    "axs_flat[2].plot(history_times_rl, history_o_points_rl[:, 2],'b.')\n",
    "axs_flat[2].plot(history_times_nl, history_o_points_nl[:, 2],'rx')\n",
    "axs_flat[2].set_xlabel('Time')\n",
    "axs_flat[2].set_ylabel('Magnetic axis, $\\psi$')\n",
    "axs_flat[2].grid()\n",
    "\n",
    "axs_flat[3].plot(history_times, history_currents[:,-1]*stepping.plasma_norm_factor,'k+')\n",
    "axs_flat[3].plot(history_times_rl, history_currents_rl[:,-1]*stepping.plasma_norm_factor,'b.')\n",
    "axs_flat[3].plot(history_times_nl, history_currents_nl[:,-1]*stepping.plasma_norm_factor,'rx')\n",
    "axs_flat[3].set_xlabel('Time')\n",
    "axs_flat[3].set_ylabel('Plasma current')\n",
    "axs_flat[3].grid()\n",
    "\n",
    "axs_flat[4].plot(history_times, history_elongation,'k+')\n",
    "axs_flat[4].plot(history_times_rl, history_elongation_rl,'b.')\n",
    "axs_flat[4].plot(history_times_nl, history_elongation_nl,'rx')\n",
    "axs_flat[4].set_xlabel('Time')\n",
    "axs_flat[4].set_ylabel('Geometric elongation')\n",
    "axs_flat[4].grid()\n",
    "\n",
    "axs_flat[5].plot(history_times, history_triangularity,'k+')\n",
    "axs_flat[5].plot(history_times_rl, history_triangularity_rl,'b.')\n",
    "axs_flat[5].plot(history_times_nl, history_triangularity_nl,'rx')\n",
    "axs_flat[5].set_xlabel('Time')\n",
    "axs_flat[5].set_ylabel('Triangularity')\n",
    "axs_flat[5].grid()\n",
    "\n",
    "axs_flat[6].plot(history_times, history_squareness,'k+')\n",
    "axs_flat[6].plot(history_times_rl, history_squareness_rl,'b.')\n",
    "axs_flat[6].plot(history_times_nl, history_squareness_nl,'rx')\n",
    "axs_flat[6].set_xlabel('Time')\n",
    "axs_flat[6].set_ylabel('Squarenss')\n",
    "axs_flat[6].grid()\n",
    "\n",
    "axs_flat[7].plot(history_times, history_area,'k+')\n",
    "axs_flat[7].plot(history_times_rl, history_area_rl,'b.')\n",
    "axs_flat[7].plot(history_times_nl, history_area_nl,'rx')\n",
    "axs_flat[7].set_xlabel('Time')\n",
    "axs_flat[7].set_ylabel('LCFS area')\n",
    "axs_flat[7].grid()\n",
    "\n",
    "axs_flat[8].plot(history_times, history_length,'k+')\n",
    "axs_flat[8].plot(history_times_rl, history_length_rl,'b.')\n",
    "axs_flat[8].plot(history_times_nl, history_length_nl,'rx')\n",
    "axs_flat[8].set_xlabel('Time')\n",
    "axs_flat[8].set_ylabel('LCFS circumference')\n",
    "axs_flat[8].grid()\n",
    "\n",
    "\n",
    "# plot relinearisation times\n",
    "relin_times = history_times_rl[0:-1][history_jtor_norm_rl > threshold]\n",
    "for i in range(0,9):\n",
    "    axs_flat[i].vlines(relin_times, ymin=[axs_flat[i].get_ylim()[0]]*len(relin_times), ymax=[axs_flat[i].get_ylim()[1]]*len(relin_times), linestyles=\"--\", color=\"k\", linewidths=0.7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the equilibria at the final time step\n",
    "fig1, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 8), dpi=60)\n",
    "\n",
    "ax1.grid(True, which='both')\n",
    "history_equilibria[-1].plot(axis=ax1, show=False)\n",
    "eq.tokamak.plot(axis=ax1, show=False)\n",
    "ax1.set_xlim(0.1, 2.15)\n",
    "ax1.set_ylim(-2.25, 2.25)\n",
    "ax1.set_title(\"Linear\")\n",
    "\n",
    "ax2.grid(True, which='both')\n",
    "history_equilibria_rl[-1].plot(axis=ax2, show=False)\n",
    "eq.tokamak.plot(axis=ax2, show=False)\n",
    "ax2.set_xlim(0.1, 2.15)\n",
    "ax2.set_ylim(-2.25, 2.25)\n",
    "ax2.set_title(\"Relinearised\")\n",
    "\n",
    "ax3.grid(True, which='both')\n",
    "history_equilibria_nl[-1].plot(axis=ax3, show=False)\n",
    "eq.tokamak.plot(axis=ax3, show=False)\n",
    "ax3.set_xlim(0.1, 2.15)\n",
    "ax3.set_ylim(-2.25, 2.25)\n",
    "ax3.set_title(\"Nonlinear\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot relinearisation criteria evolution\n",
    "fig, axs = plt.subplots(1, 1, figsize=(12, 4), dpi=80)\n",
    "axs.plot(history_times[0:-1], history_jtor_norm, 'k+-', label=\"linear\")\n",
    "axs.plot(history_times_rl[0:-1], history_jtor_norm_rl, 'b.-', label=\"relinearised\")\n",
    "axs.hlines(threshold, xmin=history_times[0], xmax=history_times[-1], linestyle='--', color='k', label=\"Threshold\")\n",
    "axs.set_xlabel('Time')\n",
    "axs.set_ylabel(r\"Relative $J_{\\phi}$ norm\")\n",
    "axs.set_yscale('log')\n",
    "axs.legend()\n",
    "axs.grid()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot runtimes\n",
    "fig, axs = plt.subplots(2, 1, figsize=(10, 6), dpi=80)\n",
    "axs = axs.flat\n",
    "\n",
    "axs[0].plot(history_times[0:-1], history_timings, 'k+-', label=\"linear\")\n",
    "axs[0].plot(history_times_rl[0:-1], history_timings_rl, 'b.-', label=\"relinearised\")\n",
    "axs[0].plot(history_times_nl[0:-1], history_timings_nl, 'rx-', label=\"nonlinear\")\n",
    "axs[0].set_ylabel('Runtime per step [s]')\n",
    "axs[0].set_yscale('log')\n",
    "axs[0].grid()\n",
    "axs[0].set_xticklabels([])\n",
    "\n",
    "axs[1].plot(history_times[0:-1], np.cumsum(history_timings)/60, 'k+-', label=\"linear\")\n",
    "axs[1].plot(history_times_rl[0:-1], np.cumsum(history_timings_rl)/60, 'b.-', label=\"relinearised\")\n",
    "axs[1].plot(history_times_nl[0:-1], np.cumsum(history_timings_nl)/60, 'rx-', label=\"nonlinear\")\n",
    "\n",
    "\n",
    "# plot relinearisation times\n",
    "for i in range(0,2):\n",
    "    axs[i].vlines(relin_times, ymin=[axs[i].get_ylim()[0]]*len(relin_times), ymax=[axs[i].get_ylim()[1]]*len(relin_times), linestyles=\"--\", color=\"k\", linewidths=0.9)\n",
    "\n",
    "# axs.plot(jtor_norm_nl)\n",
    "axs[1].set_xlabel('Simulation time [s]')\n",
    "axs[1].set_ylabel('Cumulative runtime [mins]')\n",
    "axs[1].legend()\n",
    "axs[1].grid()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Re-linearising when using the `no_GS` option\n",
    "\n",
    "Next, we will do the same simulations (with and without relinearisation) but this time without solving Grad-Shafranov at each timestep (see prior example notebook). This will enable a very fast simulation but, due to the lack of a control scheme, the accuracy of the simulation will be lost very quickly (less so when relinearisation is switched on). \n",
    "\n",
    "We should note that given GS is no longer solved, the criterion for relinearisation is no longer valid (as we will not have access to $J_{\\phi}$ at each timestep). We therefore carry out relinearisation when the absolute change in the `plasma_descriptors` (with respect to the values at the last linearisation) exceeds $\\epsilon$. \n",
    "\n",
    "A value for $\\epsilon$ can be set for each descriptor individually (as a list of floats) or for all of them (as a single float)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let us carry out a standard linear simulation without solving GS at each timestep (no relinearisation is carried out). Here we set the `relinearise_threshold` as a list of `None` (or large) values to indicate that we do not relinearise at all. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.dIydI_ICs = None\n",
    "stepping.dIydtheta_ICs = None\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "# storage\n",
    "history_times_nogs = [t]\n",
    "history_currents_nogs = [stepping.currents_vec]\n",
    "history_plasma_descriptors_nogs = [plasma_descriptors(eq)]\n",
    "history_criteria_nogs = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulate\n",
    "while counter < max_count:\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "\n",
    "    # carry out the time step (feed in the voltages and profile parameters)\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages,\n",
    "        linear_only=True,\n",
    "        no_GS=True,\n",
    "        verbose=False,\n",
    "        relinearise_threshold=[None, 1.0, 1.0, 1.0],\n",
    "    )\n",
    "\n",
    "    # store time-advanced currents and plasma descriptors\n",
    "    history_currents_nogs.append(stepping.currents_vec)\n",
    "    history_plasma_descriptors_nogs.append(stepping.plasma_descriptors_vec)\n",
    "    history_criteria_nogs.append(stepping.relinearise_criteria)\n",
    "\n",
    "    t += stepping.dt_step\n",
    "    history_times_nogs.append(t)\n",
    "    counter += 1\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_nogs = np.array(history_currents_nogs)\n",
    "history_times_nogs = np.array(history_times_nogs)\n",
    "history_plasma_descriptors_nogs = np.array(history_plasma_descriptors_nogs)\n",
    "history_criteria_nogs = np.array(history_criteria_nogs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll do the same simulation but this time we will enable re-linearisation when one of the shape parameter changes exceeds a few centimetres."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset the solver object by resetting the initial conditions\n",
    "stepping.dIydI_ICs = None\n",
    "stepping.dIydtheta_ICs = None\n",
    "stepping.initialize_from_ICs(eq, profiles)\n",
    "\n",
    "counter = 0\n",
    "t = 0\n",
    "\n",
    "# storage\n",
    "history_times_nogs_rl = [t]\n",
    "history_currents_nogs_rl = [stepping.currents_vec]\n",
    "history_plasma_descriptors_nogs_rl = [plasma_descriptors(eq)]\n",
    "history_criteria_nogs_rl = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simulate\n",
    "thresholds = [0.01, 0.02, 0.02, None]\n",
    "\n",
    "currents = []\n",
    "profs = []\n",
    "params = []\n",
    "while counter < max_count:\n",
    "    print(f'Step: {counter}/{max_count-1}')\n",
    "    print(f'--- t = {t:.2e}')\n",
    "\n",
    "    # carry out the time step (feed in the voltages and profile parameters)\n",
    "    stepping.nlstepper(\n",
    "        active_voltage_vec=voltages,\n",
    "        linear_only=True,\n",
    "        no_GS=True,\n",
    "        verbose=False,\n",
    "        relinearise_threshold=thresholds,\n",
    "    )\n",
    "\n",
    "    # store time-advanced currents and plasma descriptors\n",
    "    history_currents_nogs_rl.append(stepping.currents_vec)\n",
    "    history_plasma_descriptors_nogs_rl.append(stepping.plasma_descriptors_vec)\n",
    "    history_criteria_nogs_rl.append(stepping.relinearise_criteria)\n",
    "\n",
    "    t += stepping.dt_step\n",
    "    history_times_nogs_rl.append(t)\n",
    "    counter += 1\n",
    "\n",
    "    currents.append(stepping.initial_currents_plasma_descriptor)\n",
    "    profs.append(stepping.initial_profiles_plasma_descriptor)   \n",
    "    params.append(stepping.initial_plasma_descriptors)\n",
    "\n",
    "# transform lists to arrays\n",
    "history_currents_nogs_rl = np.array(history_currents_nogs_rl)\n",
    "history_times_nogs_rl = np.array(history_times_nogs_rl)\n",
    "history_plasma_descriptors_nogs_rl = np.array(history_plasma_descriptors_nogs_rl)\n",
    "history_criteria_nogs_rl = np.array(history_criteria_nogs_rl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(4,1, figsize=(10, 12), dpi=80)\n",
    "ax = ax.flatten()\n",
    "\n",
    "labels = [\"Zcurrent [m]\", \"Rx [m]\", \"Zx [m]\", \"Rin [m]\"]\n",
    "\n",
    "for i in range(4):\n",
    "    ax[i].plot(history_times, history_plasma_descriptors_nl[:, i], 'rx-', label=\"Nonlinear\")\n",
    "    # ax[i].plot(history_times, history_plasma_descriptors[:, i], 'k+-', label=\"Linear (with GS)\")\n",
    "    ax[i].plot(history_times_nogs, history_plasma_descriptors_nogs[:, i], 'm3-', label=\"Linear (without GS)\")\n",
    "    ax[i].plot(history_times_nogs_rl, history_plasma_descriptors_nogs_rl[:, i], 'g.-', label=\"Linear (without GS, with relinearisation)\")\n",
    "    ax[i].set_xlabel(\"Time (s)\")\n",
    "    ax[i].set_ylabel(labels[i])\n",
    "    ax[i].legend()\n",
    "    ax[i].grid()\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot relinearisation criteria evolution\n",
    "fig, ax = plt.subplots(4,1, figsize=(10, 12), dpi=80)\n",
    "ax = ax.flatten()\n",
    "\n",
    "for i in range(4):\n",
    "    ax[i].plot(history_times_nogs_rl[0:-1], history_criteria_nogs_rl[:, i], 'b.-', label=\"Linear (without GS, with relinearisation)\")\n",
    "    ax[i].plot(history_times_nogs[0:-1], history_criteria_nogs[:, i], 'k+-', label=\"Linear (without GS)\")\n",
    "    ax[i].hlines(thresholds[i], xmin=history_times_nogs_rl[0], xmax=history_times_nogs_rl[-1], linestyle='--', color='k', label=\"Threshold\")\n",
    "    ax[i].set_xlabel(\"Time (s)\")\n",
    "    ax[i].set_ylabel(f\"Relative change in {labels[i]}\")\n",
    "    ax[i].set_yscale('log')\n",
    "    ax[i].legend()\n",
    "    ax[i].grid()\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "\n",
    "# fig, axs = plt.subplots(1, 1, figsize=(12, 4), dpi=80)\n",
    "# # axs.plot(history_times_nogs[0:-1], history_jtor_norm, 'k+-', label=\"linear\")\n",
    "# axs.plot(history_times_nogs[0:-1], history_criteria_nogs, 'b.-', label=\"relinearised\")\n",
    "# axs.hlines(threshold, xmin=history_times[0], xmax=history_times[-1], linestyle='--', color='k', label=\"Threshold\")\n",
    "# axs.set_xlabel('Time')\n",
    "# axs.set_ylabel(r\"Relative $J_{\\phi}$ norm\")\n",
    "# axs.set_yscale('log')\n",
    "# axs.legend()\n",
    "# axs.grid()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(12,1, figsize=(10, 24), dpi=80)\n",
    "ax = ax.flatten()\n",
    "\n",
    "labels = eq.tokamak.coil_names[0:12]\n",
    "\n",
    "for i in range(len(labels)):\n",
    "    ax[i].plot(history_times, history_currents_nl[:, i], 'rx-', label=\"Nonlinear\")\n",
    "    # ax[i].plot(history_times, history_plasma_descriptors[:, i], 'k+-', label=\"Linear (with GS)\")\n",
    "    ax[i].plot(history_times_nogs, history_currents_nogs[:, i], 'm3-', label=\"Linear (without GS)\")\n",
    "    ax[i].plot(history_times_nogs_rl, history_currents_nogs_rl[:, i], 'g.-', label=\"Linear (without GS, with relinearisation)\")\n",
    "    ax[i].set_xlabel(\"Time (s)\")\n",
    "    ax[i].set_ylabel(labels[i])\n",
    "    ax[i].legend()\n",
    "    ax[i].grid()\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}